"""
This is part of Human Robot Dynamics-Control, By Tourists.

Based on "OpenLoong Dynamics Control"
Transfer from mujoco to ISAAC SIM by IsaacLab

"""
#pragma once

#include "qpOASES.hpp"
#include <algorithm>
#include <Eigen/Dense>
#include "data_bus.h"
#include "useful_math.h"
#include "priority_tasks.h"
#include "pino_kin_dyn.h"
#include <iostream>
#include <iomanip>

from ..common.data_bus import DataBus, MotionState, LegState
from ..mathq.useful_math import eul2Rot, diffRot
from .priority_tasks import PriorityTasks

import numpy as np
import quadprog
from scipy.optimize import minimize

# test_task = PriorityTasks()
# test_task.addTask("test_task")

# class Test_debug:
#     def __init__(self):
#         self.test_task = PriorityTasks()
#         self.test_task.addTask("test_task")
    


class WBC_priority:
    def __init__(self, model_nv_In, QP_nvIn, QP_ncIn, miu_In, dt):
        self.timeStep = dt
        self.model_nv = model_nv_In
        self.miu = miu_In
        self.nWSR = 100
        self.last_nWSR = 0
        self.cpu_time=0.1
        self.last_cpu_time = 0
        self.qpStatus = 0
        self.QP_nc = QP_ncIn
        self.QP_nv = QP_nvIn
        
        self.J_base = np.empty((0, 0), dtype=np.float64)
        self.dJ_base = np.empty((0, 0), dtype=np.float64)
        self.Jcom = np.empty((0, 0), dtype=np.float64)
        self.J_hip_link = np.empty((0, 0), dtype=np.float64)
        self.base_pos_des = np.zeros(3)
        self.base_pos = np.zeros(3)
        self.base_rpy_des = np.zeros(3)
        self.base_rpy_cur = np.zeros(3)
        self.hip_link_pos = np.zeros(3)
        self.hip_link_rot = np.zeros((3, 3), dtype=np.float64)
        self.base_rot = np.zeros((3, 3), dtype=np.float64)
        self.swing_fe_pos_des_W = np.empty(0, dtype=np.float64)
        self.swing_fe_rpy_des_W = np.empty(0, dtype=np.float64)
        self.stance_fe_pos_cur_W = np.zeros(3)
        self.stance_fe_rot_cur_W = np.zeros((3, 3), dtype=np.float64)
        self.stanceDesPos_W = np.zeros(3)
        self.des_ddq = np.empty(0, dtype=np.float64)
        self.des_dq = np.empty(0, dtype=np.float64)
        self.des_delta_q = np.empty(0, dtype=np.float64)
        self.des_q = np.empty(0, dtype=np.float64)
        self.qIniDes = np.empty(0, dtype=np.float64)
        self.qIniCur = np.empty(0, dtype=np.float64)
        
        self.tauJointRes = np.empty(0, dtype=np.float64)
        self.delta_q_final_kin = np.empty(0, dtype=np.float64)
        self.dq_final_kin = np.empty(0, dtype=np.float64)
        self.pCoMDes = np.zeros(3)
        self.pCoMCur = np.zeros(3)
        self.hip_link_rot = np.zeros((3, 3), dtype=np.float64)
        self.base_rot = np.zeros((3, 3), dtype=np.float64)
        
        self.l_shoulder_pitch = 0.0 ## q(28) - qIniDes(28);
        self.r_shoulder_pitch = 0.0 ## q(34) - qIniDes(34);
        # self.ddq_final_kin = np.array()
        
        self.motionStateCur = MotionState.Stand
        
        self.qIniDes = np.empty(0, dtype=np.float64)
        self.qIniCur = np.empty(0, dtype=np.float64)
        
        self.qpStatus = 0
        
        self.Sf = np.zeros((6, model_nv_In))
        self.Sf[:6, :6] = np.eye(6)
        
        self.St_qpV2 = np.zeros((model_nv_In, model_nv_In - 6))
        self.St_qpV2[6:, :model_nv_In - 6] = np.eye(model_nv_In - 6)
        
        self.St_qpV1 = np.zeros((model_nv_In, 6))
        self.St_qpV1[:6, :6] = np.eye(6)
        
        self.f_z_low = 10
        self.f_z_upp = 1400
        
        self.tau_upp_stand_L = np.array([10, 20, 40])
        self.tau_low_stand_L = np.array([-10, -20, -40])
        
        self.tau_upp_walk_L = np.array([25, 40, 40])
        self.tau_low_walk_L = np.array([-25, -40, -40])
        
        self.eigen_xOpt = np.zeros(self.QP_nv)
        self.eigen_ddq_Opt = np.zeros(self.model_nv)
        self.eigen_fr_Opt = np.zeros(12)
        self.eigen_tau_Opt = np.zeros(self.model_nv - 6)
        
        self.delta_q_final_kin = np.zeros(self.model_nv)
        self.dq_final_kin = np.zeros(self.model_nv)
        self.ddq_final_kin = np.zeros(self.model_nv)
        
        self.base_rpy_cur = np.zeros(3)
        
        # Initialize tasks
        self.kin_tasks_walk = PriorityTasks()
        self.kin_tasks_walk.addTask("static_Contact")
        self.kin_tasks_walk.addTask("Roll_Pitch_Yaw_Pz")
        self.kin_tasks_walk.addTask("RedundantJoints")
        self.kin_tasks_walk.addTask("PxPy")
        self.kin_tasks_walk.addTask("SwingLeg")
        self.kin_tasks_walk.addTask("HandTrack")
        self.kin_tasks_walk.addTask("HandTrackJoints")
        self.kin_tasks_walk.addTask("PosRot")
        
        taskOrder_walk = ["RedundantJoints", "static_Contact", "PosRot", "SwingLeg", "HandTrackJoints"]
        self.kin_tasks_walk.buildPriority(taskOrder_walk)
        
        self.kin_tasks_stand = PriorityTasks()
        self.kin_tasks_stand.addTask("static_Contact")
        self.kin_tasks_stand.addTask("CoMTrack")
        self.kin_tasks_stand.addTask("HandTrackJoints")
        self.kin_tasks_stand.addTask("HipRPY")
        self.kin_tasks_stand.addTask("HeadRP")
        self.kin_tasks_stand.addTask("Pz")
        self.kin_tasks_stand.addTask("CoMXY_HipRPY")
        self.kin_tasks_stand.addTask("Roll_Pitch_Yaw")
        self.kin_tasks_stand.addTask("fixedWaist")
        
        taskOrder_stand = ["static_Contact", "CoMXY_HipRPY", "Pz", "HandTrackJoints", "HeadRP"]
        self.kin_tasks_stand.buildPriority(taskOrder_stand)

    def dataBusRead(self, robot_state):
        # foot-end offset posture
        self.fe_L_rot_L_off = robot_state.fe_L_rot_L_off
        self.fe_R_rot_L_off = robot_state.fe_R_rot_L_off

        # desired values
        self.base_rpy_des = robot_state.base_rpy_des
        self.base_rpy_cur = np.array([robot_state.rpy[0], robot_state.rpy[1], robot_state.rpy[2]])
        self.base_pos_des = robot_state.base_pos_des
        self.swing_fe_pos_des_W = robot_state.swing_fe_pos_des_W
        self.swing_fe_rpy_des_W = robot_state.swing_fe_rpy_des_W
        self.stance_fe_pos_cur_W = robot_state.stance_fe_pos_cur_W
        self.stance_fe_rot_cur_W = robot_state.stance_fe_rot_cur_W
        self.stanceDesPos_W = robot_state.stanceDesPos_W
        self.hd_l_pos_cur_W = robot_state.hd_l_pos_W
        self.hd_r_pos_cur_W = robot_state.hd_r_pos_W
        self.hd_l_rot_cur_W = robot_state.hd_l_rot_W
        self.hd_r_rot_cur_W = robot_state.hd_r_rot_W
        self.fe_l_pos_cur_W = robot_state.fe_l_pos_W
        self.fe_r_pos_cur_W = robot_state.fe_r_pos_W
        self.fe_l_rot_cur_W = robot_state.fe_l_rot_W
        self.fe_r_rot_cur_W = robot_state.fe_r_rot_W
        self.des_ddq = robot_state.des_ddq
        self.des_dq = robot_state.des_dq
        self.des_delta_q = robot_state.des_delta_q
        self.des_q = robot_state.des_q

        # state update
        self.J_base = robot_state.J_base
        self.dJ_base = robot_state.dJ_base
        self.base_rot = robot_state.base_rot
        self.base_pos = robot_state.base_pos
        self.hip_link_pos = robot_state.hip_link_pos
        self.hip_link_rot = robot_state.hip_link_rot
        self.J_hip_link = robot_state.J_hip_link

        self.Jfe = np.zeros((12, self.model_nv))
        self.Jfe[:6, :] = robot_state.J_l
        self.Jfe[6:, :] = robot_state.J_r
        self.dJfe = np.zeros((12, self.model_nv))
        self.dJfe[:6, :] = robot_state.dJ_l
        self.dJfe[6:, :] = robot_state.dJ_r
        self.J_hd_l = robot_state.J_hd_l
        self.J_hd_r = robot_state.J_hd_r
        self.dJ_hd_l = robot_state.J_hd_l
        self.dJ_hd_r = robot_state.J_hd_r
        self.Fr_ff = robot_state.Fr_ff
        self.dyn_M = robot_state.dyn_M
        self.dyn_M_inv = robot_state.dyn_M_inv
        self.dyn_Ag = robot_state.dyn_Ag
        self.dyn_dAg = robot_state.dyn_dAg
        self.dyn_Non = robot_state.dyn_Non
        self.dq = robot_state.dq
        self.q = robot_state.q
        self.legStateCur = robot_state.legState
        self.motionStateCur = robot_state.motionState

        if self.legStateCur == LegState.LSt:
            self.Jc = robot_state.J_l
            self.dJc = robot_state.dJ_l
            self.Jsw = robot_state.J_r
            self.dJsw = robot_state.dJ_r
            self.fe_pos_sw_W = robot_state.fe_r_pos_W
            self.fe_rot_sw_W = robot_state.fe_r_rot_W
        else:
            self.Jc = robot_state.J_r
            self.dJc = robot_state.dJ_r
            self.Jsw = robot_state.J_l
            self.dJsw = robot_state.dJ_l
            self.fe_pos_sw_W = robot_state.fe_l_pos_W
            self.fe_rot_sw_W = robot_state.fe_l_rot_W

        self.Jcom = robot_state.Jcom_W
        self.pCoMCur = robot_state.pCoM_W

    def dataBusWrite(self, robot_state):
        robot_state.wbc_ddq_final = self.eigen_ddq_Opt
        robot_state.wbc_tauJointRes = self.tauJointRes
        robot_state.wbc_FrRes = self.eigen_fr_Opt
        robot_state.qp_cpuTime = self.cpu_time
        robot_state.qp_nWSR = self.nWSR
        robot_state.qp_status = self.qpStatus

        robot_state.wbc_delta_q_final = self.delta_q_final_kin
        robot_state.wbc_dq_final = self.dq_final_kin
        robot_state.wbc_ddq_final = self.ddq_final_kin

        robot_state.qp_status = self.qpStatus
        robot_state.qp_nWSR = self.nWSR
        robot_state.qp_cpuTime = self.cpu_time

    ## QP problem contains joint torque, QP_nv=6+12, QP_nc=22;
    def computeTau(self):
     # 构建等式约束矩阵
        eigen_qp_A1 = np.zeros((6, self.QP_nv))
        eigen_qp_A1[:6, :6] = self.Sf @ self.dyn_M @ self.St_qpV1
        eigen_qp_A1[:6, 6:] = -self.Sf @ self.Jfe.T

        # 设置等式约束结果
        eqRes = -self.Sf @ (self.dyn_M @ self.ddq_final_kin + self.dyn_Non) + self.Sf @ (self.Jfe.T @ self.Fr_ff)

        # 计算旋转矩阵
        if self.motionStateCur == 'Stand':
            Rfe = self.fe_l_rot_cur_W
        else:
            Rfe = self.stance_fe_rot_cur_W

        # 构建 Mw2b 矩阵
        Mw2b = np.zeros((12, 12))
        Mw2b[:3, :3] = Rfe.T
        Mw2b[3:6, 3:6] = Rfe.T
        Mw2b[6:9, 6:9] = Rfe.T
        Mw2b[9:, 9:] = Rfe.T

        # 构建 W 矩阵
        W = np.zeros((16, 12))
        W[0, 0] = 1
        W[0, 2] = np.sqrt(2) / 2.0 * self.miu
        W[1, 0] = -1
        W[1, 2] = np.sqrt(2) / 2.0 * self.miu
        W[2, 1] = 1
        W[2, 2] = np.sqrt(2) / 2.0 * self.miu
        W[3, 1] = -1
        W[3, 2] = np.sqrt(2) / 2.0 * self.miu
        W[4:8, 2:6] = np.eye(4)
        W[8:, 6:] = W[:8, :6]
        W = W @ Mw2b

        # 设置上下限
        f_low = np.zeros(16)
        f_upp = np.zeros(16)
        tau_upp_fe = self.tau_upp_stand_L if self.motionStateCur == 'Stand' else self.tau_upp_walk_L
        tau_low_fe = self.tau_low_stand_L if self.motionStateCur == 'Stand' else self.tau_low_walk_L

        f_upp[:8] = [1e10, 1e10, 1e10, 1e10, self.f_z_upp, tau_upp_fe[0], tau_upp_fe[1], tau_upp_fe[2]]
        f_upp[8:] = f_upp[:8]
        f_low[:8] = [0, 0, 0, 0, self.f_z_low, tau_low_fe[0], tau_low_fe[1], tau_low_fe[2]]
        f_low[8:] = f_low[:8]

        if self.motionStateCur in ['Walk', 'Walk2Stand']:
            if self.legStateCur == 'LSt':
                f_upp[12:] = 0
                f_low[12:] = 0
                f_low[8:12] = -1e-7
            elif self.legStateCur == 'RSt':
                f_upp[4:8] = 0
                f_low[4:8] = 0
                f_low[:4] = -1e-7

        # 构建不等式约束矩阵
        eigen_qp_A2 = np.zeros((16, 18))
        eigen_qp_A2[:, 6:] = W
        neqRes_low = f_low - W @ self.Fr_ff
        neqRes_upp = f_upp - W @ self.Fr_ff

        # 构建最终约束矩阵
        eigen_qp_A_final = np.zeros((self.QP_nc, self.QP_nv))
        eigen_qp_A_final[:6, :] = eigen_qp_A1
        eigen_qp_A_final[6:, :] = eigen_qp_A2

        # 设置等式和不等式约束结果
        eigen_qp_lbA = np.zeros(22)
        eigen_qp_ubA = np.zeros(22)
        eigen_qp_lbA[:6] = eqRes
        eigen_qp_lbA[6:] = neqRes_low
        eigen_qp_ubA[:6] = eqRes
        eigen_qp_ubA[6:] = neqRes_upp

        # 构建 H 矩阵
        eigen_qp_H = np.zeros((self.QP_nv, self.QP_nv))
        Q2 = np.eye(6)
        Q1 = np.eye(12)
        eigen_qp_H[:6, :6] = Q2 * 2.0 * 1e7
        eigen_qp_H[6:, 6:] = Q1 * 2.0 * 1e1

        # 定义目标函数
        def objective(x):
            return 0.5 * x.T @ eigen_qp_H @ x

        # 定义约束条件
        constraints = [
            {'type': 'eq', 'fun': lambda x: eqRes - eigen_qp_A1 @ x},
            {'type': 'ineq', 'fun': lambda x: neqRes_low - eigen_qp_A2 @ x},
            {'type': 'ineq', 'fun': lambda x: eigen_qp_A2 @ x - neqRes_upp}
        ]

        # 初始猜测值
        x0 = np.zeros(self.QP_nv)

        # 求解优化问题
        result = minimize(objective, x0, constraints=constraints)
        # x_opt, _, _, _ = quadprog.solve_qp(eigen_qp_H, g, A_ineq, b_ineq, A_eq=A_eq, b_eq=b_eq)

        # 获取最优解
        eigen_xOpt = result.x
        # eigen_xOpt = x_opt

        # 更新 ddq 和 Fr
        eigen_ddq_Opt = self.ddq_final_kin.copy()
        eigen_ddq_Opt[:6] += eigen_xOpt[:6]
        eigen_fr_Opt = self.Fr_ff + eigen_xOpt[6:]

        # 计算 tau
        tauRes = self.dyn_M @ eigen_ddq_Opt + self.dyn_Non - self.Jfe.T @ eigen_fr_Opt
        self.tauJointRes = tauRes[6:]

        # return tauJointRes
        self.last_nWSR = self.nWSR
        self.last_cpu_time = self.cpu_time

    def computeDdq(self, pinKinDynIn):
        # // task definition
        # /// -------- walk -------------
        # static_Contact 任务
        id = self.kin_tasks_walk.getId("static_Contact")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(3)
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(3)
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(3)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(3)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(3) * 0
        self.kin_tasks_walk.taskLib[id].kd = np.eye(3) * 0
        self.kin_tasks_walk.taskLib[id].J = self.Jc[:3, :self.model_nv]
        self.kin_tasks_walk.taskLib[id].dJ = self.dJc[:3, :self.model_nv]
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)

        # RedundantJoints 任务
        id = self.kin_tasks_walk.getId("RedundantJoints")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(5)
        self.kin_tasks_walk.taskLib[id].errX[0] = 0 - self.q[21]
        self.kin_tasks_walk.taskLib[id].errX[1] = 0 - self.q[22]
        self.kin_tasks_walk.taskLib[id].errX[2] = 0 - self.q[23]
        self.kin_tasks_walk.taskLib[id].errX[3] = 0 - self.q[24]
        self.kin_tasks_walk.taskLib[id].errX[4] = 0 - self.q[25]
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(5)
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(5)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(5)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(5) * 200
        self.kin_tasks_walk.taskLib[id].kd = np.eye(5) * 20
        self.kin_tasks_walk.taskLib[id].J = np.zeros((5, self.model_nv))
        self.kin_tasks_walk.taskLib[id].J[0, 20] = 1
        self.kin_tasks_walk.taskLib[id].J[1, 21] = 1
        self.kin_tasks_walk.taskLib[id].J[2, 22] = 1
        self.kin_tasks_walk.taskLib[id].J[3, 23] = 1
        self.kin_tasks_walk.taskLib[id].J[4, 24] = 1
        self.kin_tasks_walk.taskLib[id].dJ = np.zeros((5, self.model_nv))
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)

        # Roll_Pitch_Yaw_Pz 任务
        id = self.kin_tasks_walk.getId("Roll_Pitch_Yaw_Pz")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(4)
        desRot = eul2Rot(self.base_rpy_des[0], self.base_rpy_des[1], self.base_rpy_des[2])
        self.kin_tasks_walk.taskLib[id].errX[:3] = diffRot(self.base_rot, desRot)
        self.kin_tasks_walk.taskLib[id].errX[3] = self.base_pos_des[2] - self.q[2]
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(4)
        self.kin_tasks_walk.taskLib[id].derrX[:3] = -self.dq[3:6]
        self.kin_tasks_walk.taskLib[id].derrX[3] = 0 - self.dq[2]
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(4)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(4)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(4) * 2000
        self.kin_tasks_walk.taskLib[id].kd = np.eye(4) * 100
        taskMap = np.zeros((4, 6))
        taskMap[0, 3] = 1
        taskMap[1, 4] = 1
        taskMap[2, 5] = 1
        taskMap[3, 2] = 1
        self.kin_tasks_walk.taskLib[id].J = taskMap @ self.J_base
        self.kin_tasks_walk.taskLib[id].dJ = taskMap @ self.dJ_base
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)

        # PxPy 任务
        id = self.kin_tasks_walk.getId("PxPy")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(2)
        # self.kin_tasks_walk.taskLib[id].errX = self.des_dq[:2, 0] * self.timeStep
        self.kin_tasks_walk.taskLib[id].errX = self.des_dq[:2] * self.timeStep
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(2)
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(2)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(2)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(2) * 500
        self.kin_tasks_walk.taskLib[id].kd = np.eye(2) * 50
        taskMap = np.zeros((2, 6))
        taskMap[0, 0] = 1
        taskMap[1, 1] = 1
        self.kin_tasks_walk.taskLib[id].J = taskMap @ self.J_base
        self.kin_tasks_walk.taskLib[id].dJ = taskMap @ self.dJ_base
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)

        # PosRot 任务
        id = self.kin_tasks_walk.getId("PosRot")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].errX[:3] = self.base_pos_des[:3] - self.q[:3]
        if abs(self.kin_tasks_walk.taskLib[id].errX[0]) >= 0.02:
            self.kin_tasks_walk.taskLib[id].errX[0] = 0.02 * np.sign(self.kin_tasks_walk.taskLib[id].errX[0])
        if abs(self.kin_tasks_walk.taskLib[id].errX[1]) >= 0.01:
            self.kin_tasks_walk.taskLib[id].errX[1] = 0.01 * np.sign(self.kin_tasks_walk.taskLib[id].errX[1])
        desRot = eul2Rot(self.base_rpy_des[0], self.base_rpy_des[1], self.base_rpy_des[2])
        self.kin_tasks_walk.taskLib[id].errX[3:] = diffRot(self.base_rot, desRot)
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(6)
        self.kin_tasks_walk.taskLib[id].kp[3:, 3:] = np.eye(3) * 2000
        self.kin_tasks_walk.taskLib[id].kd = np.eye(6) * 2
        self.kin_tasks_walk.taskLib[id].kd[3:, 3:] = np.eye(3) * 100
        self.kin_tasks_walk.taskLib[id].J = self.J_base
        self.kin_tasks_walk.taskLib[id].dJ = self.dJ_base
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)

        # SwingLeg 任务
        id = self.kin_tasks_walk.getId("SwingLeg")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].errX[:3] = self.swing_fe_pos_des_W - self.fe_pos_sw_W
        desRot = eul2Rot(self.swing_fe_rpy_des_W[0], self.swing_fe_rpy_des_W[1], self.swing_fe_rpy_des_W[2])
        self.kin_tasks_walk.taskLib[id].errX[3:] = diffRot(self.fe_rot_sw_W, desRot)
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(6)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(6) * 2000
        self.kin_tasks_walk.taskLib[id].kp[2, 2] = self.kin_tasks_walk.taskLib[id].kp[2, 2] * 0.1
        self.kin_tasks_walk.taskLib[id].kd = np.eye(6) * 20
        self.kin_tasks_walk.taskLib[id].J = self.Jsw
        self.kin_tasks_walk.taskLib[id].J[:, 22:25] = 0  # 排除腰部关节
        self.kin_tasks_walk.taskLib[id].dJ = self.dJsw
        self.kin_tasks_walk.taskLib[id].dJ[:, 22:25] = 0  # 排除腰部关节
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)
        
        hd_l_eul_L_des = np.array([-1.7581, 0.2129, 2.9581])
        hd_r_eul_L_des = np.array([1.7581, 0.2129, -2.9581])

        hd_l_rot_des = eul2Rot(hd_l_eul_L_des[0], hd_l_eul_L_des[1], hd_l_eul_L_des[2])
        hd_r_rot_des = eul2Rot(hd_r_eul_L_des[0], hd_r_eul_L_des[1], hd_r_eul_L_des[2])

        base2shoulder_l_pos_L_des = np.array([0.0040, 0.1616, 0.3922])
        shoulder2hand_l_pos_L_des = np.array([-0.0240, 0.1584, -0.5512])
        base2shoulder_r_pos_L_des = np.array([0.0040, -0.1616, 0.3922])
        shoulder2hand_r_pos_L_des = np.array([-0.0240, -0.1584, -0.5512])

        l_hip_pitch = self.q[28] - self.qIniDes[28]
        r_hip_pitch = self.q[34] - self.qIniDes[34]
        k = 0.8

        hd_l_rot_des = eul2Rot(0, -k * r_hip_pitch, 0) @ hd_l_rot_des
        hd_r_rot_des = eul2Rot(0, -k * l_hip_pitch, 0) @ hd_r_rot_des

        hd_l_pos_W_des = (
            eul2Rot(0, -k * r_hip_pitch, 0) @ shoulder2hand_l_pos_L_des
            + base2shoulder_l_pos_L_des + self.base_pos
        )
        hd_r_pos_W_des = (
            eul2Rot(0, -k * l_hip_pitch, 0) @ shoulder2hand_r_pos_L_des
            + base2shoulder_r_pos_L_des + self.base_pos
        )

        hd_l_pos_L_des = (
            eul2Rot(0, -k * r_hip_pitch, 0) @ shoulder2hand_l_pos_L_des
            + base2shoulder_l_pos_L_des
        )
        hd_r_pos_L_des = (
            eul2Rot(0, -k * l_hip_pitch, 0) @ shoulder2hand_r_pos_L_des
            + base2shoulder_r_pos_L_des
        )

        hd_l_rot_W_des = hd_l_rot_des
        hd_r_rot_W_des = hd_r_rot_des

        # HandTrack 任务
        id = self.kin_tasks_walk.getId("HandTrack")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(12)
        self.kin_tasks_walk.taskLib[id].errX[:3] = hd_l_pos_W_des - self.hd_l_pos_cur_W
        self.kin_tasks_walk.taskLib[id].errX[3:6] = diffRot(self.hd_l_rot_cur_W, hd_l_rot_W_des)
        self.kin_tasks_walk.taskLib[id].errX[6:9] = hd_r_pos_W_des - self.hd_r_pos_cur_W
        self.kin_tasks_walk.taskLib[id].errX[9:] = diffRot(self.hd_r_rot_cur_W, hd_r_rot_W_des)
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(12)
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(12)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(12)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(12) * 2000
        self.kin_tasks_walk.taskLib[id].kd = np.eye(12) * 20
        self.kin_tasks_walk.taskLib[id].J = np.zeros((12, self.model_nv))
        self.kin_tasks_walk.taskLib[id].J[:6, :] = self.J_hd_l
        self.kin_tasks_walk.taskLib[id].J[6:, :] = self.J_hd_r
        self.kin_tasks_walk.taskLib[id].dJ = np.zeros((12, self.model_nv))
        self.kin_tasks_walk.taskLib[id].dJ[:6, :] = self.dJ_hd_l
        self.kin_tasks_walk.taskLib[id].dJ[6:, :] = self.dJ_hd_r
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)

        # 计算 resLeg
        resLeg = pinKinDynIn.computeInK_Hand(hd_l_rot_des, hd_l_pos_L_des, hd_r_rot_des, hd_r_pos_L_des)

        # HandTrackJoints 任务
        id = self.kin_tasks_walk.getId("HandTrackJoints")
        self.kin_tasks_walk.taskLib[id].errX = np.zeros(14)
        self.kin_tasks_walk.taskLib[id].errX = resLeg.jointPosRes[:14] - self.q[7:21]
        self.kin_tasks_walk.taskLib[id].derrX = np.zeros(14)
        self.kin_tasks_walk.taskLib[id].ddxDes = np.zeros(14)
        self.kin_tasks_walk.taskLib[id].dxDes = np.zeros(14)
        self.kin_tasks_walk.taskLib[id].kp = np.eye(14) * 2000
        self.kin_tasks_walk.taskLib[id].kd = np.eye(14) * 100
        self.kin_tasks_walk.taskLib[id].J = np.zeros((14, self.model_nv))
        self.kin_tasks_walk.taskLib[id].J[:14, 6:20] = np.eye(14)
        self.kin_tasks_walk.taskLib[id].dJ = np.zeros((14, self.model_nv))
        self.kin_tasks_walk.taskLib[id].W = np.eye(self.model_nv)

        # /// -------- stand -------------
        # # 获取任务 ID
        id = self.kin_tasks_stand.getId("static_Contact")

        # 初始化静态接触任务
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(12)
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(12)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(12)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(12)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(12) * 0
        self.kin_tasks_stand.taskLib[id].kd = np.eye(12) * 0
        self.kin_tasks_stand.taskLib[id].J = np.zeros((12, self.model_nv))
        taskCtMap = np.zeros((3, 3))
        taskCtMap[1, 1] = 1
        taskCtMap[2, 2] = 1
        taskCtMap = self.fe_l_rot_cur_W @ taskCtMap @ self.fe_l_rot_cur_W.T
        self.kin_tasks_stand.taskLib[id].J = self.Jfe
        self.kin_tasks_stand.taskLib[id].J[3:6, :] = taskCtMap @ self.kin_tasks_stand.taskLib[id].J[3:6, :]
        self.kin_tasks_stand.taskLib[id].J[9:12, :] = taskCtMap @ self.kin_tasks_stand.taskLib[id].J[9:12, :]
        self.kin_tasks_stand.taskLib[id].J[:, 22:] = 0
        self.kin_tasks_stand.taskLib[id].dJ = np.zeros((12, self.model_nv))
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        # 初始化髋部姿态控制任务
        id = self.kin_tasks_stand.getId("HipRPY")
        desRot = eul2Rot(0, 0, 0)
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].errX[:3] = diffRot(self.hip_link_rot, desRot)
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(3) * 1000
        self.kin_tasks_stand.taskLib[id].kd = np.eye(3) * 50
        taskMapRPY = np.zeros((3, 6))
        taskMapRPY[0, 3] = 1
        taskMapRPY[1, 4] = 1
        taskMapRPY[2, 5] = 1
        self.kin_tasks_stand.taskLib[id].J = taskMapRPY @ self.J_hip_link
        self.kin_tasks_stand.taskLib[id].J[:, 22:25] = 0
        self.kin_tasks_stand.taskLib[id].J[:, 6:20] = 0
        self.kin_tasks_stand.taskLib[id].dJ = np.zeros((3, self.model_nv))
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        # 初始化基座高度控制任务
        id = self.kin_tasks_stand.getId("Pz")
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(1)
        self.kin_tasks_stand.taskLib[id].errX[0] = self.base_pos_des[2] - self.q[2]
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(1)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(1)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(1)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(1) * 2000
        self.kin_tasks_stand.taskLib[id].kd = np.eye(1) * 10
        taskMap = np.zeros((1, 6))
        taskMap[0, 2] = 1
        self.kin_tasks_stand.taskLib[id].J = taskMap @ self.J_base
        self.kin_tasks_stand.taskLib[id].J[:, 22:25] = 0
        self.kin_tasks_stand.taskLib[id].dJ = taskMap @ self.dJ_base
        self.kin_tasks_stand.taskLib[id].dJ[:, 22:25] = 0
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        # 初始化质心位置跟踪任务
        id = self.kin_tasks_stand.getId("CoMTrack")
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].errX = self.pCoMDes[:2] - self.pCoMCur[:2]
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(2) * 2000
        self.kin_tasks_stand.taskLib[id].kd = np.eye(2) * 100
        self.kin_tasks_stand.taskLib[id].J = self.Jcom[:2, :]
        self.kin_tasks_stand.taskLib[id].J[:, 6:20] = 0
        self.kin_tasks_stand.taskLib[id].dJ = np.zeros((2, self.model_nv))
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        # 初始化质心位置与髋部姿态复合控制任务
        id = self.kin_tasks_stand.getId("CoMXY_HipRPY")
        taskMapRPY = np.zeros((3, 6))
        taskMapRPY[0, 3] = 1
        taskMapRPY[1, 4] = 1
        taskMapRPY[2, 5] = 1
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(5)
        self.kin_tasks_stand.taskLib[id].errX[:2] = self.pCoMDes[:2] - self.pCoMCur[:2]
        desRot = eul2Rot(self.base_rpy_des[0], self.base_rpy_des[1], self.base_rpy_des[2])
        self.kin_tasks_stand.taskLib[id].errX[2:5] = diffRot(self.hip_link_rot, desRot)
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(5)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(5)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(5)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(5) * 1000
        self.kin_tasks_stand.taskLib[id].kd = np.eye(5) * 10
        self.kin_tasks_stand.taskLib[id].kp[2:5, 2:5] = np.eye(3) * 1000
        self.kin_tasks_stand.taskLib[id].kd[2:5, 2:5] = np.eye(3) * 10
        self.kin_tasks_stand.taskLib[id].J = np.zeros((5, self.model_nv))
        self.kin_tasks_stand.taskLib[id].J[:2, :] = self.Jcom[:2, :]
        self.kin_tasks_stand.taskLib[id].J[2:5, :] = taskMapRPY @ self.J_hip_link
        self.kin_tasks_stand.taskLib[id].J[2:5, 22:25] = 0
        self.kin_tasks_stand.taskLib[id].J[2:5, 6:20] = 0
        self.kin_tasks_stand.taskLib[id].dJ = np.zeros((5, self.model_nv))
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)
        self.kin_tasks_stand.taskLib[id].W[22] = 200
        self.kin_tasks_stand.taskLib[id].W[23] = 200

        # 定义摆动臂运动
        hd_l_eul_L_des = np.array([-1.7581, 0.2129, 2.9581])
        hd_r_eul_L_des = np.array([1.7581, 0.2129, -2.9581])
        hd_l_rot_des = eul2Rot(hd_l_eul_L_des[0], hd_l_eul_L_des[1], hd_l_eul_L_des[2])
        hd_r_rot_des = eul2Rot(hd_r_eul_L_des[0], hd_r_eul_L_des[1], hd_r_eul_L_des[2])

        base2shoulder_l_pos_L_des = np.array([0.0040, 0.1616, 0.3922])
        shoulder2hand_l_pos_L_des = np.array([-0.0240, 0.1584, -0.5512])
        base2shoulder_r_pos_L_des = np.array([0.0040, -0.1616, 0.3922])
        shoulder2hand_r_pos_L_des = np.array([-0.0240, -0.1584, -0.5512])

        k = 1
        hd_l_rot_des = eul2Rot(0, -k * self.r_shoulder_pitch, 0) @ hd_l_rot_des
        hd_r_rot_des = eul2Rot(0, -k * self.l_shoulder_pitch, 0) @ hd_r_rot_des

        hd_l_pos_L_des = eul2Rot(0, -k * self.r_shoulder_pitch, 0) @ shoulder2hand_l_pos_L_des + base2shoulder_l_pos_L_des
        hd_r_pos_L_des = eul2Rot(0, -k * self.l_shoulder_pitch, 0) @ shoulder2hand_r_pos_L_des + base2shoulder_r_pos_L_des

        resLeg = pinKinDynIn.computeInK_Hand(hd_l_rot_des, hd_l_pos_L_des, hd_r_rot_des, hd_r_pos_L_des)

        # 初始化手部关节跟踪任务
        id = self.kin_tasks_stand.getId("HandTrackJoints")
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(14)
        self.kin_tasks_stand.taskLib[id].errX = resLeg.jointPosRes[:14] - self.q[7:21]
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(14)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(14)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(14)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(14) * 2000
        self.kin_tasks_stand.taskLib[id].kd = np.eye(14) * 100
        self.kin_tasks_stand.taskLib[id].J = np.zeros((14, self.model_nv))
        self.kin_tasks_stand.taskLib[id].J[:14, 6:20] = np.eye(14)
        self.kin_tasks_stand.taskLib[id].dJ = np.zeros((14, self.model_nv))
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        # 初始化头部姿态控制任务
        id = self.kin_tasks_stand.getId("HeadRP")
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].errX[0] = 0 - self.q[21]
        self.kin_tasks_stand.taskLib[id].errX[1] = self.base_rpy_cur[1] - self.q[22]
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(2)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(2) * 100
        self.kin_tasks_stand.taskLib[id].kd = np.eye(2) * 10
        self.kin_tasks_stand.taskLib[id].J = np.zeros((2, self.model_nv))
        self.kin_tasks_stand.taskLib[id].J[0, 20] = 1
        self.kin_tasks_stand.taskLib[id].J[1, 21] = 1
        self.kin_tasks_stand.taskLib[id].dJ = np.zeros((2, self.model_nv))
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        # 初始化基座姿态控制任务
        id = self.kin_tasks_stand.getId("Roll_Pitch_Yaw")
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(3)
        desRot = eul2Rot(self.base_rpy_des[0], self.base_rpy_des[1], self.base_rpy_des[2])
        self.kin_tasks_stand.taskLib[id].errX = diffRot(self.base_rot, desRot)
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].derrX[:3] = -self.dq[3:6]
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].kp = np.eye(3) * 2000
        self.kin_tasks_stand.taskLib[id].kd = np.eye(3) * 100
        taskMap = np.zeros((3, 6))
        taskMap[0, 3] = 1
        taskMap[1, 4] = 1
        taskMap[2, 5] = 1
        self.kin_tasks_stand.taskLib[id].J = taskMap @ self.J_base
        self.kin_tasks_stand.taskLib[id].dJ = taskMap @ self.dJ_base
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        id = self.kin_tasks_stand.getId("fixedWaist")
        self.kin_tasks_stand.taskLib[id].errX = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].errX[0] = 0 - self.q[23]
        self.kin_tasks_stand.taskLib[id].errX[1] = 0 - self.q[24]
        self.kin_tasks_stand.taskLib[id].errX[2] = 0 - self.q[25]

        # 初始化导数和期望导数
        self.kin_tasks_stand.taskLib[id].derrX = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].ddxDes = np.zeros(3)
        self.kin_tasks_stand.taskLib[id].dxDes = np.zeros(3)

        # 设置增益矩阵
        self.kin_tasks_stand.taskLib[id].kp = np.eye(3) * 200
        self.kin_tasks_stand.taskLib[id].kd = np.eye(3) * 20

        # 初始化雅可比矩阵
        self.kin_tasks_stand.taskLib[id].J = np.zeros((3, self.model_nv))
        self.kin_tasks_stand.taskLib[id].J[0, 22] = 1
        self.kin_tasks_stand.taskLib[id].J[1, 23] = 1
        self.kin_tasks_stand.taskLib[id].J[2, 24] = 1

        # 初始化雅可比矩阵的导数
        self.kin_tasks_stand.taskLib[id].dJ = np.zeros((3, self.model_nv))

        # 设置权重矩阵
        self.kin_tasks_stand.taskLib[id].W = np.eye(self.model_nv)

        # 根据运动状态选择计算方式
        if self.motionStateCur == MotionState.Walk or self.motionStateCur == MotionState.Walk2Stand:
            self.kin_tasks_walk.computeAll(self.des_delta_q, self.des_dq, self.des_ddq, self.dyn_M, self.dyn_M_inv, self.dq)
            self.delta_q_final_kindelta_q_final_kin = self.kin_tasks_walk.out_delta_q
            self.delta_q_final_kindq_final_kin = self.kin_tasks_walk.out_dq
            self.delta_q_final_kinddq_final_kin = self.kin_tasks_walk.out_ddq
        elif self.motionStateCur == MotionState.Stand:
            self.kin_tasks_stand.computeAll(self.des_delta_q, self.des_dq, self.des_ddq, self.dyn_M, self.dyn_M_inv, self.dq)
            self.delta_q_final_kindelta_q_final_kin = self.kin_tasks_stand.out_delta_q
            self.delta_q_final_kindq_final_kin = self.kin_tasks_stand.out_dq
            self.delta_q_final_kinddq_final_kin = self.kin_tasks_stand.out_ddq
        else:
            self.delta_q_final_kindelta_q_final_kin = np.zeros(self.model_nv)
            self.delta_q_final_kindq_final_kin = np.zeros(self.model_nv)
            self.delta_q_final_kinddq_final_kin = np.zeros(self.model_nv)

        # // final WBC output collection

    def copy_eigen_to_real_t(self, target, source, n_rows, n_cols):
        count = 0

        for i in range(n_rows):
            for j in range(n_cols):
                value = source[i, j]
                if np.isinf(value):
                    target[count] = np.inf
                else:
                    target[count] = value
                count += 1

    def set_qini(self, q_ini_des_in, q_ini_cur_in):
        # global qIniDes, qIniCur
        self.qIniDes = q_ini_des_in
        self.qIniCur = q_ini_cur_in
